import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
import pytorch3d.io
from pytorch3d.structures import Meshes
from pytorch3d.renderer import (
    FoVPerspectiveCameras,
    MeshRenderer,
    MeshRasterizer,
    SoftPhongShader,
    TexturesVertex,
    RasterizationSettings,
    PointLights
)
import matplotlib.pyplot as plt
import numpy as np

# ------------------------------
# 虚拟数据集定义（需替换真实数据）
# ------------------------------
class MultiViewDataset(Dataset):
    def __init__(self, num_views=10):
        self.num_views = num_views
        # 示例：生成随机相机参数
        self.camera_params = [
            {'elev': 20+i*5, 'azim': i*36, 'dist': 0.5} 
            for i in range(num_views)
        ]
        
    def __len__(self):
        return self.num_views
    
    def __getitem__(self, idx):
        # 示例：返回随机图像和掩码（真实数据需替换为实际加载逻辑）
        return {
            'image': torch.rand(512, 512, 3),      # 假设RGB图像
            'mask': torch.ones(512, 512).bool(),   # 前景掩码
            'cam_id': idx                          # 相机标识
        }

# ------------------------------
# 可微分渲染系统
# ------------------------------
class ColorOptimizationSystem(nn.Module):
    def __init__(self, mesh, num_cameras):
        super().__init__()
        # 网格模型
        self.mesh = mesh
        # 顶点颜色参数
        self.vertex_colors = nn.Parameter(
            torch.rand(mesh._verts_packed.shape, 3)
        )
        # 相机颜色校正参数
        self.color_transforms = nn.Parameter(
            torch.eye(3)[None].repeat(num_cameras, 1, 1)  # 初始化为单位矩阵
        )
        self.color_biases = nn.Parameter(
            torch.zeros(num_cameras, 3)  # 初始无偏置
        )
        # 初始化渲染器
        self._init_renderer()
    
    def _init_renderer(self):
        raster_settings = RasterizationSettings(
            image_size=512,
            blur_radius=0.0,
            faces_per_pixel=1,
        )
        self.renderer = MeshRenderer(
            rasterizer=MeshRasterizer(raster_settings=raster_settings),
            shader=SoftPhongShader()
        )
    
    def forward(self, camera_params, cam_id):
        # 更新网格纹理
        self.mesh.textures = TexturesVertex(
            verts_features=self.vertex_colors[None]
        )
        
        # 创建对应相机
        R, T = FoVPerspectiveCameras.look_at_view_transform(
            dist=camera_params['dist'],
            elev=camera_params['elev'],
            azim=camera_params['azim'],
        )
        cameras = FoVPerspectiveCameras(R=R, T=T, device='cpu')
        
        # 执行渲染
        rendered = self.renderer(self.mesh, cameras=cameras)
        
        # 应用颜色校正
        transform = self.color_transforms[cam_id]  # [3,3]
        bias = self.color_biases[cam_id]           # 
        corrected = torch.matmul(rendered[..., :3], transform) + bias
        
        return corrected

# ------------------------------
# 损失函数
# ------------------------------
def compute_loss(pred, target, mask, model, cam_id):
    # 光度一致性损失
    photo_loss = ((pred - target)**2 * mask.unsqueeze(-1)).mean()
    
    # 颜色一致性正则
    color_consistency = model.vertex_colors.std(dim=0).mean() * 0.1
    
    # 相机参数正则
    transform_reg = model.color_transforms.norm(dim=(1,2)).mean()
    bias_reg = model.color_biases.norm(dim=1).mean()
    
    return (
        photo_loss + 
        color_consistency + 
        0.01 * (transform_reg + bias_reg)
    )

# ------------------------------
# 主流程
# ------------------------------
def main():
    # 1. 准备模型
    bunny_mesh = pytorch3d.io.load_objs_as_meshes(["./data/bunny.obj"])
    
    # 2. 初始化系统
    num_cameras = 10
    system = ColorOptimizationSystem(bunny_mesh, num_cameras)
    
    # 3. 数据加载
    dataset = MultiViewDataset(num_views=num_cameras)
    dataloader = DataLoader(dataset, batch_size=1, shuffle=True)
    
    # 4. 优化设置
    optimizer = optim.Adam([
        {'params': system.vertex_colors, 'lr': 0.05},
        {'params': [system.color_transforms, system.color_biases], 'lr': 0.01}
    ])
    
    # 5. 训练循环
    for epoch in range(100):
        total_loss = 0
        for batch in dataloader:
            # 获取数据
            cam_id = batch['cam_id'].item()
            camera_param = dataset.camera_params[cam_id]
            
            # 前向传播
            pred = system(camera_param, cam_id)
            
            # 计算损失
            loss = compute_loss(
                pred, batch['image'], batch['mask'], 
                system, cam_id
            )
            
            # 反向传播
            optimizer.zero_grad()
            loss.backward()
            
            # 梯度裁剪
            torch.nn.utils.clip_grad_norm_(system.parameters(), 0.1)
            
            optimizer.step()
            total_loss += loss.item()
        
        print(f"Epoch {epoch+1}, Loss: {total_loss/len(dataloader):.4f}")
    
    # 6. 可视化结果
    with torch.no_grad():
        test_view = dataset.camera_params
        output = system(test_view, 0)
        plt.imshow(output.numpy())
        plt.axis('off')
        plt.show()

if __name__ == "__main__":
    main()
